import asyncio
from typing import Callable, Coroutine
from datetime import datetime
from common.logger import logger

class AsyncTaskExecutor:
    def __init__(self, max_workers: int = 4, check_interval: int = 5):
        self.max_workers = max_workers
        self.check_interval = check_interval
        self.running = False
        self.tasks_queue = asyncio.Queue()
        self.workers = []

    async def start(self):
        """启动任务执行器"""
        self.running = True
        # 启动检查器和工作协程
        self.workers = [
            asyncio.create_task(self._worker(f"worker-{i}")) 
            for i in range(self.max_workers)
        ]
        asyncio.create_task(self._task_checker())
        logger.info("任务执行器已启动")

    async def stop(self):
        """停止任务执行器"""
        self.running = False
        # 等待所有任务完成
        await self.tasks_queue.join()
        # 取消所有工作协程
        for worker in self.workers:
            worker.cancel()
        await asyncio.gather(*self.workers, return_exceptions=True)
        logger.info("任务执行器已停止")

    async def add_task(self, task: Callable[..., Coroutine], *args, **kwargs):
        """添加任务到队列"""
        await self.tasks_queue.put((task, args, kwargs))
        logger.info(f"新任务已添加: {task.__name__}")

    async def _task_checker(self):
        """任务检查协程"""
        while self.running:
            try:
                await asyncio.sleep(self.check_interval)
                logger.info(f"当前队列任务数: {self.tasks_queue.qsize()}")
            except Exception as e:
                logger.error(f"任务检查出错: {str(e)}")

    async def _worker(self, name: str):
        """工作协程"""
        logger.info(f"{name} 启动")
        while self.running:
            try:
                task, args, kwargs = await self.tasks_queue.get()
                logger.info(f"{name} 开始执行任务: {task.__name__}")
                
                try:
                    result = await task(*args, **kwargs)
                    logger.info(f"{name} 任务完成: {result}")
                except Exception as e:
                    logger.error(f"{name} 任务执行失败: {str(e)}")
                finally:
                    self.tasks_queue.task_done()
                    
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"{name} 发生错误: {str(e)}")
                await asyncio.sleep(1)

# 使用示例
if __name__ == "__main__":
    async def example_task(task_id: int):
        logger.info(f"开始执行任务 {task_id}")
        await asyncio.sleep(2)  # 模拟异步耗时操作
        return f"任务 {task_id} 执行完成"

    async def main():
        # 创建任务执行器实例
        executor = AsyncTaskExecutor(max_workers=2, check_interval=3)
        await executor.start()

        # 添加一些测试任务
        for i in range(5):
            await executor.add_task(example_task, i)

        # 等待一段时间后停止
        await asyncio.sleep(20)
        await executor.stop()

    # 运行主协程
    asyncio.run(main())